{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W11_Tutorial1",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CIS-522/course-content/blob/week11_feedback/tutorials/W11_DeepRL/W11_Tutorial1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DIazMknk1W6P"
      },
      "source": [
        "# CIS-522 Week 11 Part 1\n",
        "# Introduction to Reinforcement Learning\n",
        "\n",
        "__Instructor:__ Dinesh Jayaraman\n",
        "\n",
        "__Content creators:__ Chuning Zhu\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "a5NQBmTeSs5S"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "my_pennkey = '' #@param {type:\"string\"}\n",
        "my_pod = 'Select' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','astute-jellyfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion']\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0LMCVKg3EOSn"
      },
      "source": [
        "## Recap the experience from last week\n",
        "\n",
        "What did you learn last week. What questions do you have? [10 min discussion]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0q2tL2Q5CAde",
        "cellView": "form"
      },
      "source": [
        "learning_from_previous_week = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XfcCTgt0QMnm"
      },
      "source": [
        "---\n",
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yHBnw4Nah8JZ"
      },
      "source": [
        "# imports\n",
        "\n",
        "import math\n",
        "import numpy as np\n",
        "import IPython\n",
        "from numbers import Number\n",
        "from matplotlib import pyplot as plt\n",
        "import matplotlib.patches as patches\n",
        "from tqdm.auto import tqdm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iNi1AkDhh5Lm",
        "cellView": "form"
      },
      "source": [
        "# @title Plotting functions\n",
        "%matplotlib inline \n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/\"\n",
        "              \"course-content/master/nma.mplstyle\")\n",
        "\n",
        "# Plotting functions\n",
        "def plot_episode_rewards(episode_rewards):\n",
        "    fig = plt.figure()\n",
        "    plt.plot(episode_rewards)\n",
        "    plt.xlabel(\"Episode\")\n",
        "    plt.ylabel(\"Reward\")\n",
        "    fig.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xFnuAj4MQ7sg",
        "cellView": "form"
      },
      "source": [
        "# @title Gridworld Environment\n",
        "\n",
        "'''\n",
        "    A custom Gridworld environment with deterministic transitions. Adapted from \n",
        "    CS 188 Gridworld env. There are four actions: up, left, down, right. The \n",
        "    state is the (x, y) coordinates in the Grid. \n",
        "'''\n",
        "class Gridworld():\n",
        "    \n",
        "    def __init__(self, grid, living_reward=-1.0):\n",
        "        self.h = len(grid)\n",
        "        self.w = len(grid[0])\n",
        "        self.living_reward = living_reward\n",
        "        \n",
        "        self.scale = math.ceil(max(self.h, self.w) / min(self.h, self.w))\n",
        "        self.action_space = ['up', 'left', 'down', 'right']\n",
        "        self.n_actions = 4\n",
        "\n",
        "        self.init_grid(grid)\n",
        "\n",
        "    def init_grid(self, grid):\n",
        "        # Create reward grid. The reward grid is a numpy array storing the\n",
        "        # reward given for entering each state.\n",
        "        self.rew_grid = np.array([[self.living_reward if isinstance(e, str) else e \n",
        "                                   for e in row] for row in grid], dtype=np.float)\n",
        "\n",
        "        # Create grid. The grid is a numpy of chars. \n",
        "        # S (start), T (terminal), C (cliff), # (block), or ' ' (regular).\n",
        "        convert_fn = lambda e: 'T' if e >= self.living_reward else 'C'\n",
        "        self.grid = np.array([[convert_fn(e) if isinstance(e, Number) else e \n",
        "                               for e in row] for row in grid])\n",
        "        \n",
        "        # Find initial state\n",
        "        start_indices = np.argwhere(self.grid == 'S')\n",
        "        if len(start_indices) == 0:\n",
        "            raise Exception('Grid has no start state')\n",
        "        self.init_state = (start_indices[0][1], start_indices[0][0])\n",
        "\n",
        "        \n",
        "    def get_transition(self, state, action):\n",
        "        '''\n",
        "        Execute one action in the environment.\n",
        "        Args:\n",
        "            state (tuple): the (x, y) coordinates of the current state.\n",
        "            action (int): the current action chosen from {0, 1, 2, 3}.\n",
        "\n",
        "        Returns:\n",
        "            next_state (tuple): the (x, y) coordinates of the next state.\n",
        "            reward (float): the reward for the current time step.\n",
        "        '''\n",
        "        # Handle terminal states\n",
        "        x, y = state\n",
        "        if self.grid[y, x] == 'T':\n",
        "            return state, 0\n",
        "        \n",
        "        # Handle invalid actions\n",
        "        if action not in range(len(self.action_space)):\n",
        "            raise Exception('Illegal action')\n",
        "\n",
        "        # Default transitions\n",
        "        named_action = self.action_space[action]\n",
        "        nx, ny = x, y\n",
        "        if named_action == 'up':\n",
        "            ny -= 1\n",
        "        elif named_action == 'left':\n",
        "            nx -= 1\n",
        "        elif named_action == 'down':\n",
        "            ny += 1 \n",
        "        elif named_action == 'right':\n",
        "            nx += 1\n",
        "\n",
        "        # Handle special cases\n",
        "        if nx < 0 or nx >= self.w or ny < 0 or ny >= self.h or self.grid[ny, nx] == '#':\n",
        "            # Give living reward if next state is blocked or out of bounds\n",
        "            reward = self.living_reward\n",
        "            next_state = (x, y)\n",
        "        else:\n",
        "            reward = self.rew_grid[ny, nx]\n",
        "            if self.grid[ny, nx] == 'C':\n",
        "                next_state = self.init_state # falls off cliff\n",
        "            else:\n",
        "                next_state = (nx, ny) # transition to next state\n",
        "\n",
        "        return next_state, reward\n",
        "\n",
        "    def __render(self):\n",
        "        # Render grid with matplotlib patches. \n",
        "        fig, ax = plt.subplots(figsize=(self.h*self.scale, self.w*self.scale))\n",
        "        ax.set_aspect('equal')\n",
        "        ax.set_xlim(0, self.w)\n",
        "        ax.set_ylim(0, self.h)\n",
        "        ax.set_xticklabels([])\n",
        "        ax.set_yticklabels([])\n",
        "        ax.tick_params(length=0)\n",
        "        plt.axis('off')\n",
        "        for y in range(self.h):\n",
        "            for x in range(self.w):\n",
        "                cell_type = self.grid[y, x]\n",
        "                if cell_type == 'S':\n",
        "                    c = '#DAE8FC' # blue\n",
        "                elif cell_type == '#':\n",
        "                    c = '#CCCCCC' # gray\n",
        "                elif cell_type == 'T':\n",
        "                    c = '#D5E8D4' # green\n",
        "                elif cell_type == 'C':\n",
        "                    c = '#F8CECC' # red\n",
        "                else:\n",
        "                    c = '#FFFFFF' # white\n",
        "                rect = patches.Rectangle((x, self.h-y-1), 1, 1, fc=c, ec='gray', lw=1)\n",
        "                ax.add_patch(rect)\n",
        "        return fig, ax\n",
        "\n",
        "    def render_grid(self):\n",
        "        fig, ax = self.__render()\n",
        "        for y in range(self.h):\n",
        "            for x in range(self.w):\n",
        "                if self.grid[y, x] != '#':\n",
        "                    # alternate: x+0.1, self.h-y-0.2\n",
        "                    ax.text(x+0.5, self.h-y-0.5, str(self.rew_grid[y, x]), size='medium', ha='center', va='center')\n",
        "        plt.title(\"Rewards\")\n",
        "        fig.show()\n",
        "    \n",
        "    def render_values(self, V):\n",
        "        fig, ax = self.__render()\n",
        "        for y in range(self.h):\n",
        "            for x in range(self.w):\n",
        "                ax.text(x+0.5, self.h-y-0.5, '{:.2f}'.format(V[y, x]), size='medium', ha='center', va='center')\n",
        "        plt.title(\"Values\")\n",
        "        fig.show()\n",
        "\n",
        "    def render_q_values(self, Q):\n",
        "        fig, ax = self.__render()\n",
        "        for y in range(self.h):\n",
        "            for x in range(self.w):\n",
        "                named_action = self.action_space[np.argmax(Q[y, x])]\n",
        "                xl, xc, xr = x, x+0.5, x+1\n",
        "                yt, yc, yb = self.h-y, self.h-y-0.5, self.h-y-1\n",
        "                ce, tl, bl, tr, br = [xc, yc], [xl, yt], [xl, yb], [xr, yt], [xr, yb]\n",
        "                if named_action == 'up':\n",
        "                    xy = np.array([ce, tl, tr])\n",
        "                elif named_action == 'left':\n",
        "                    xy = np.array([ce, tl, bl])\n",
        "                elif named_action == 'down':\n",
        "                    xy = np.array([ce, bl, br])\n",
        "                elif named_action == 'right':\n",
        "                    xy = np.array([ce, br, tr])\n",
        "                ax.plot([x, x+1], [self.h-y, self.h-y-1], 'gray', lw=1)\n",
        "                ax.plot([x, x+1], [self.h-y-1, self.h-y], 'gray', lw=1)\n",
        "                poly = patches.Polygon(xy, True, fc='#FFFF00', ec='gray')\n",
        "                ax.add_patch(poly)\n",
        "                ax.text(x+0.5, self.h-y-0.2, '{:.2f}'.format(Q[y, x, 0]), size='small', ha='center', va='center')\n",
        "                ax.text(x+0.2, self.h-y-0.5, '{:.2f}'.format(Q[y, x, 1]), size='small', ha='center', va='center')\n",
        "                ax.text(x+0.5, self.h-y-0.8, '{:.2f}'.format(Q[y, x, 2]), size='small', ha='center', va='center')\n",
        "                ax.text(x+0.8, self.h-y-0.5, '{:.2f}'.format(Q[y, x, 3]), size='small', ha='center', va='center')\n",
        "        fig.show()\n",
        "        plt.title(\"Q-values\")\n",
        "        pass\n",
        "    \n",
        "    def render_policy(self, policy):\n",
        "        fig, ax = self.__render()\n",
        "        for y in range(self.h):\n",
        "            for x in range(self.w):\n",
        "                if policy[y, x] not in range(len(self.action_space)):\n",
        "                    raise Exception('Illegal action')\n",
        "                if self.grid[y, x] == 'T':\n",
        "                    continue\n",
        "                arrow_len = 0.3\n",
        "                dx, dy = 0, 0\n",
        "                named_action = self.action_space[policy[y, x]]\n",
        "                if named_action == 'up':\n",
        "                    dy = arrow_len\n",
        "                elif named_action == 'left':\n",
        "                    dx = -arrow_len\n",
        "                elif named_action == 'down':\n",
        "                    dy = -arrow_len\n",
        "                elif named_action == 'right':\n",
        "                    dx = arrow_len\n",
        "                arrow = patches.FancyArrow(x+0.5, self.h-y-0.5, dx, dy, 0.03, True, color='#6C8EBF')\n",
        "                ax.add_patch(arrow)\n",
        "        plt.title(\"Policy\")\n",
        "        fig.show()\n",
        "\n",
        "\n",
        "'''\n",
        "    GridworldEnv is a wrapper around Gridworld implementing an RL interface.\n",
        "'''\n",
        "class GridworldEnv(Gridworld):\n",
        "    def __init__(self, grid, living_reward=-1.0):\n",
        "        super().__init__(grid, living_reward)\n",
        "        self.reset()\n",
        "        \n",
        "    def reset(self):\n",
        "        '''\n",
        "        Reset the agent to its initial state\n",
        "        '''\n",
        "        self.state = self.init_state\n",
        "        return self.state\n",
        "\n",
        "    def step(self, action):\n",
        "        '''\n",
        "        Execute one action in the environment.\n",
        "        Args:\n",
        "            action (int): the current action chosen from {0, 1, 2, 3}.\n",
        "\n",
        "        Returns:\n",
        "            next_state (tuple): (x, y) coordinates of the next state.\n",
        "            reward (float): reward for the current time step.\n",
        "            done (bool): True if a terminal state has been reached, False otherwise.\n",
        "        '''\n",
        "        next_state, reward = self.get_transition(self.state, action)\n",
        "        self.state = next_state\n",
        "        done = self.grid[self.state[1], self.state[0]] == 'T'\n",
        "        return next_state, reward, done\n",
        "\n",
        "# Pre-defined grids\n",
        "def get_book_grid():\n",
        "    grid = [['T',  ' ',  ' ',  ' '],\n",
        "            [' ',  ' ',  ' ',  ' '],\n",
        "            [' ',  ' ',  ' ',  ' '], \n",
        "            ['S',  ' ',  ' ',  'T']]\n",
        "    return GridworldEnv(grid)\n",
        "\n",
        "def get_cliff_small():\n",
        "    grid = [[' ',   ' ',  ' ',  ' ',  ' '],\n",
        "            ['S',   ' ',  ' ',  ' ',   'T'],\n",
        "            [-100, -100, -100, -100, -100]]\n",
        "    return GridworldEnv(grid)\n",
        "\n",
        "def get_cliff_walk():\n",
        "    grid = [[' ' for _ in range(12)] for _ in range(3)]\n",
        "    grid.append([-100 for _ in range(12)])\n",
        "    grid[3][ 0] = 'S'\n",
        "    grid[3][-1] = 'T'\n",
        "    return GridworldEnv(grid)\n",
        "\n",
        "def get_bridge_grid():\n",
        "    grid = [[ '#',-100, -100, -100, -100, -100, '#'],\n",
        "            [   1, 'S',  ' ',  ' ',  ' ',  ' ',  10],\n",
        "            [ '#',-100, -100, -100, -100, -100, '#']]\n",
        "    return GridworldEnv(grid)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QrwdQ2uZ326v"
      },
      "source": [
        "---\n",
        "# Section 1: Introduction\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BSJWWkvq0xUv",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Intro to Reinforcement Learning\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"cVTud58UfpQ\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tOvIv7bd0ssI"
      },
      "source": [
        "Up to this point, we have mainly been concerned with supervised learning. In a supervised learning problem, we are provided with a dataset where each sample comes with a ground truth label (e.g. class label), and the goal is to learn to predict the label by minimizing some loss function. Reinforcement learning, on the other hand, is a framework for solving sequential decision-making problems. Consider an agent operating in some environment. The agent's goal is to carry out the best sequence of actions that maximizes the cumulative reward. This is difficult because the action at the current time step influences future states of the environment, which then feed back to the agent's observations. The following figure illustrates this setting. \n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/CIS-522/course-content/main/tutorials/W11_DeepRL/static/agent_environment.png\" alt=\"RL setting\" />\n",
        "\n",
        "\n",
        "What is the role of reinforcement learning in intelligence? According to Yann LeCun, if intelligence is a cake, then unsupervised learning is the bulk of the cake, supervised learning the icing, and reinforcement learning the cherry on top. The reason RL takes up such a small proportion is that very little learning in real world comes with explicit reward signal. This analogy is still debatable, as some RL folks argue that intelligence is more like a cake with lots of cherries on top, especially after the invention of [hindight experience replay](https://arxiv.org/abs/1707.01495). In addition, there are ways to solve sequential decision making problems without relying on shaped rewards, such as inverse reinforcement learning, which infers a reward function from experience, and learning from goals / demonstrations / examples. Another way to put RL in perspective is by comparing it with vision and natural language processing. If we decompose intelligence into perception, cognition (reasoning), and action (decision making), then vision coarsely corresponds to perception, NLP cognition, and RL action. Just like how vision can be combined with NLP for tasks like image captioning, RL can be organically combined with vision and NLP as well. \n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/CIS-522/course-content/main/tutorials/W11_DeepRL/static/cherry_cream_cake.jpeg\" alt=\"Cherry cake\" />\n",
        "\n",
        "In this first tutorial, we will briefly step away from deep learning and study a few classic approaches in reinforcement learning. A good reference is Sutton and Barto's book, Reinforcement Learning: An Introduction. The [full text](http://incompleteideas.net/book/the-book.html) is avaliable online."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "89WTR6ApXHvk"
      },
      "source": [
        "---\n",
        "# Section 2: MDP and Bellman Equations\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TMGCAwaxfOta"
      },
      "source": [
        "## Section 2.1: Markov Decision Process"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qx4b3ejy1OD9",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Markov Decision Processes\n",
        "try: t1;\n",
        "except NameError: t1=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"GJEL-QkT2yk\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fw9tTf82fA-4"
      },
      "source": [
        "\n",
        "We begin our study of reinforcement learning with a definition of  Markov decision process. A Markov decision process (MDP) is a tuple $(S, A, P, R, \\gamma)$, where\n",
        "\n",
        "- $S$ is the set of **states**.\n",
        "- $A$ is the set of **actions**.\n",
        "- $P$ defines the **transition probablities**. $P(s'|s, a)$ gives the probability of transitioning to state $s'$ by taking action $a$ at state $s$. \n",
        "- $R$ is the **reward function**. $R(s, a)$ gives the reward of taking action $a$ at state $s$. $R$ can also be a function of state only.\n",
        "- $\\gamma$ is the **discount factor**. It controls how much future rewards matter to us. We will talk more about discount factor in the next video.\n",
        "\n",
        "As an aside, we introduce partially observable MDP (POMDP). A POMDP additionally has a set of obervations $O$ and emission probabilities $\\varepsilon$. $\\varepsilon(o|s)$ gives the probability of observing $o$ at state $s$. This formulation is useful when we don't have access to explicit state information, but are provided with observations that may not fully reveal the underlying states. An example is reinforcement learning from images.\n",
        "\n",
        "Come up with a sequential decision making problem and formalize it as an MDP. What are $S$, $A$, $P$, and $R$? Share your example with your pod."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39V-FrCpWb-2",
        "cellView": "form"
      },
      "source": [
        "MDP_example = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hi-ZCWHcfSGS"
      },
      "source": [
        "## Section 2.2 Solving MDPs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vqUOUj4B29aT",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Solving MDPs\n",
        "\n",
        "video = YouTubeVideo(id=\"meywaLPitZ4\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O16_8EVc21Jz"
      },
      "source": [
        "\n",
        "A policy $\\pi$ is a function mapping states to distributions over actions. At state $s$, we sample an action from the distribution $\\pi(a|s)$ to execute in the environment. If all probability mass is assigned to one action, then the policy is deterministic. The goal of reinforcement learning is to find an optimal policy that maximizes the expected sum of discounted rewards: $$E\\left[\\sum_{t=0}^{\\infty}\\gamma^tR(s_t, a_t)\\right]$$\n",
        "\n",
        "Note that this objective assumes a continuous task, i.e. that $t$ extends to infinity. We can generalize it to episodic tasks with finite horizons by replacing $\\infty$ with task horizon $T$. We may also discard the discount factor $\\gamma$ in an episodic task.\n",
        "\n",
        "Before we move on to some heavy math, consider this interesting (and somewhat philosophical) question: does life has a discount factor? Why or why not?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i21694GJh2zP",
        "cellView": "form"
      },
      "source": [
        "life_discount = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4XxYIO5H4WAk"
      },
      "source": [
        "## Section 2.3: Bellman Equations\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x1eA0Uc22tRT",
        "cellView": "form"
      },
      "source": [
        "#@title Video : V, Q, and the Bellman Equation\n",
        "try: t2;\n",
        "except NameError: t2=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"tm39P5jT320\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pLGIBaui4IfK"
      },
      "source": [
        "We define the value of a state $s$ under policy $\\pi$ as the expected future reward for following $\\pi$ starting from $s$: $$V^{\\pi}(s) = E_{\\pi} \\left[\\sum_{t'=t}^{\\infty} \\gamma^{t'-t}R(s_{t'}, a_{t'}) \\mid s_t = s\\right] $$\n",
        "\n",
        "We further define the value of a state-action pair $(s, a)$ under policy $\\pi$ as the expected future reward for taking action $a$ at state $s$, *and then* following $\\pi$. This is also known as the Q-value. $$Q^{\\pi}(s, a) = E_{\\pi} \\left[\\sum_{t'=t}^{\\infty} \\gamma^{t'-t}R(s_{t'}, a_{t'}) \\mid s_t = s, a_t = a\\right]$$\n",
        "\n",
        "Observe that $V$ and $Q$ can be related by a simple equation:\n",
        "$$V^{\\pi}(s) = E_{a\\sim \\pi(a|s)}\\left[Q^{\\pi}(s, a)\\right]$$\n",
        "\n",
        "By definition, $V$ and $Q$ satisfy the following Bellman equations.\n",
        "\n",
        "\\begin{align*}\n",
        "V^{\\pi}(s) &= E_{a \\sim \\pi(a|s)}\\left[R(s, a)+ \\gamma E_{s' \\sim P(s'|s, a)} \\left[V^{\\pi}(s')\\right]\\right] \\\\\n",
        "Q^{\\pi}(s, a) &= R(s, a) + \\gamma E_{s' \\sim P(s'|s, a)}\\left[ E_{a' \\sim \\pi(a'|s')}\\left[Q(s', a')\\right]\\right]\n",
        "\\end{align*}\n",
        "\n",
        "The optimal value function capture the expected future reward if we start from state $s$ and act optimally in the future. Similarly, the optimal Q-function captures the expected future reward if we start from state $s$, take action $a$, and then act optimally in the future. They satisfy the Bellman optimality equations:   \n",
        "\n",
        "\\begin{align*}\n",
        "V^*(s) &= \\max_{a\\in A}\\left(R(s, a) + \\gamma E_{s' \\sim P(s'|s, a)} \\left[V^*(s')\\right]\\right)\\\\\n",
        "Q^*(s, a) &= R(s, a) + \\gamma E_{s' \\sim P(s'|s, a)} \\left[ \\max_{a' \\in A} Q^*(s', a')\\right]\n",
        "\\end{align*}\n",
        "\n",
        "If we have learned the optimal value function $V^*$ or Q-function $Q^*$, we can infer an optimal (deterministic) policy known as the greedy policy or argmax policy: $$\\pi(s) = \\arg\\max_{a\\in A}Q^*(s, a)$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gjpQmkho4CMK"
      },
      "source": [
        "---\n",
        "# Section 3: Gridworld\n",
        "\n",
        "\n",
        "Before we dive into RL algorithms, let's get familiar with the running example we will use throughout this tutorial -- the Gridworld environment. As its name suggests, the Gridworld environment is an $m \\times n$ grid. The states are x-y coordinates in the grid, with origin at the top-left corner. The agent starts from the initial state and aims to reach the terminal state. There are four actions: up, left, down, and right. Each action leads to a **deterministic** transition to the adjacent cell in the correspond direction. By default, a reward of -1 is issued for entering any non-terminal state, although our implementation allows you to define an arbitrary reward for each state. \n",
        "\n",
        "To be more specific, our implementation admits four special cells: 'S' (start), 'T' (terminal), 'C' (cliff), and '#' (block). They are colored blue, green, red, and gray respectively. The agent spawns at the start cell, and aims to reach the terminal cell. If the agent \"falls off the cliff,\" it will get a high penalty (-100) and be sent back to the start cell. If the agent tries to enter a block cell or go out of the grid, it will instead stay at the same place and get a reward of -1. \n",
        "\n",
        "Familiarize yourself with the environment's interface by interacting with the following code cells. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qVmIOoR2BtIV"
      },
      "source": [
        "# Get a pre-defined grid\n",
        "gw = get_cliff_small()\n",
        "\n",
        "# Render rewards \n",
        "gw.render_grid()\n",
        "\n",
        "# Render random values\n",
        "values = np.random.rand(4, 12)\n",
        "gw.render_values(values)\n",
        "\n",
        "# Render random Q-values and argmax policy\n",
        "q_values = np.random.randn(4, 12, 4)\n",
        "gw.render_q_values(q_values)\n",
        "\n",
        "# Render random policy\n",
        "policy = np.random.choice(4, (4, 12)).astype(int)\n",
        "gw.render_policy(policy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcY20RwdmF68"
      },
      "source": [
        "In our Gridworld environment, states are represented by a tuple (x, y), and actions are encoded as 0, 1, 2, 3, corresponding to up, left, down, right. `reset()` resets the agent to its initial state and returns the initial state. `step(action)` executes an action in the environment. It returns the agent's next state, the reward, and a boolean value indicating whether or not the terminal state is reached. In the following cell, control the agent to reach the terminal state."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Su9q2dZcDXIa"
      },
      "source": [
        "action_space = ['up', 'left', 'down', 'right']\n",
        "def gw_step(gw, a):\n",
        "    next_state, reward, done = gw.step(a)\n",
        "    print(f'You moved {action_space[a]} to {next_state}, reward: {reward}, terminal state reached: {done}')\n",
        "\n",
        "print(f\"Initial state: {gw.reset()}\") # reset to initial state\n",
        "gw_step(gw, 0) # move up\n",
        "gw_step(gw, 2) # move down\n",
        "\n",
        "# Use gw_step() to reach the terminal state.\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o60dzYCbmUQ0"
      },
      "source": [
        "A useful method of the `Gridworld` class is `get_transition(state, action)`. It takes in a state and an action and returns the next state and the reward. We will use this function for exercises 1-3 where we assume full knowledge of the environment's transitions. In a reinforcement learning setting, we only have access to `step(action)`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-q013KaKmTx_"
      },
      "source": [
        "# Show next state and reward for each action at state (0, 1)\n",
        "print(gw.get_transition((0, 1), 0))\n",
        "print(gw.get_transition((0, 1), 1))\n",
        "print(gw.get_transition((0, 1), 2))\n",
        "print(gw.get_transition((0, 1), 3))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_IbGOBZ0XW4Q"
      },
      "source": [
        "---\n",
        "# Section 4 Dynamic Programming"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5XUdhM8f4UyS",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Policy and Value Iteration\n",
        "try: t3;\n",
        "except NameError: t3=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"l87rgLg90HI\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ic_SJemA4C85"
      },
      "source": [
        "## Section 4.1: Policy Iteration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cThkPW0pN81H"
      },
      "source": [
        "If we have full knowledge of the environment, in particular its transitions, we can use dynamic programming to find the optimal policy. The first algorithm we will study is policy iteration. We start with policy evaluation, which computes the value function of the policy using the Bellman equation. We iteratively perform Bellman backup for the value of each state until convergence: \n",
        "$$V(s) \\leftarrow \\sum_{a} \\pi(a|s) \\left(R(s, a) + \\gamma\\sum_{s'}P(s'|s, a)V(s')\\right) $$\n",
        "Since we have deterministic transitions, this simplifies to \n",
        "$$V(s) \\leftarrow \\sum_{a} \\pi(a|s) \\left(R(s, a) + \\gamma V(s')\\right)$$\n",
        "where $s'$ is the state we transition to by taking action $a$ at state $s$. \n",
        "\n",
        "In the following excercise, you will evaluate a random policy which assigns equal probablities to all actions at each state. Complete one step of Bellman backup. You can get the next state and reward using `grid.get_transition((x, y), action)`.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ska_I22ua1jd"
      },
      "source": [
        "### Exercise 1"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2GOVgQpKRhzL"
      },
      "source": [
        "# Random Policy evaluation\n",
        "def random_policy_evaluation(grid, gamma=1.0):\n",
        "    values = np.zeros_like(grid.rew_grid)\n",
        "    iter = 0\n",
        "    while True:\n",
        "        eps = 0\n",
        "        for y in range(grid.h):\n",
        "            for x in range(grid.w):\n",
        "                v = values[y, x]\n",
        "                new_v = 0\n",
        "                for action in range(grid.n_actions):\n",
        "                     ###########################################################\n",
        "                     # Fill in missing code below (...),\n",
        "                     # then remove or comment the line below to test your function\n",
        "                     raise NotImplementedError(\"Random policy evaluation\")\n",
        "                     ###########################################################\n",
        "                     new_v += ...\n",
        "                     \n",
        "                values[y, x] = new_v\n",
        "                eps = max(eps, abs(new_v - v))\n",
        "        iter += 1\n",
        "        if eps < 0.0001:\n",
        "            print(\"Converged after {} iterations\".format(iter))\n",
        "            break\n",
        "    return values\n",
        "\n",
        "# # Uncomment to test\n",
        "# grid = get_book_grid()\n",
        "# values = random_policy_evaluation(grid)\n",
        "# grid.render_values(values)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ijdT36NGFq8t"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "# Random Policy evaluation\n",
        "def random_policy_evaluation(grid, gamma=1.0):\n",
        "    values = np.zeros_like(grid.rew_grid)\n",
        "    iter = 0\n",
        "    while True:\n",
        "        eps = 0\n",
        "        for y in range(grid.h):\n",
        "            for x in range(grid.w):\n",
        "                v = values[y, x]\n",
        "                new_v = 0\n",
        "                for action in range(grid.n_actions):\n",
        "                    (new_x, new_y), reward = grid.get_transition((x, y), action)\n",
        "                    new_v += 0.25 * (reward + gamma * values[new_y, new_x])\n",
        "                values[y, x] = new_v\n",
        "                eps = max(eps, abs(new_v - v))\n",
        "        iter += 1\n",
        "        if eps < 0.0001:\n",
        "            print(\"Converged after {} iterations\".format(iter))\n",
        "            break\n",
        "    return values\n",
        "\n",
        "grid = get_book_grid()\n",
        "values = random_policy_evaluation(grid)\n",
        "grid.render_values(values)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qBpGwHYwE5-Z"
      },
      "source": [
        "Now we move on to the policy iteration algorithm. Policy iteration consists of two steps: policy evaluation and policy improvement. We first evaluate the policy, and then use the new values to derive a better policy by selecting the greedy action at each state. These steps are repeated until convergence. For an analysis of the theoretical guarantees of policy iteration, see [this page](http://incompleteideas.net/book/first/ebook/node42.html).\n",
        "\n",
        "In the following exercise, you will implement the policy iteration algorithm. For policy evaluation, note that we have a deterministic greedy policy, so there's no need to iterate over actions. The general backup rule becomes $$V(s) \\leftarrow R(s, \\pi(s)) + \\gamma\\sum_{s'} P(s'|s, \\pi(s)) V(s')$$ and for our environment this is $V(s) \\leftarrow R(s, \\pi(s)) + \\gamma V(s')$. For policy improvement, we do the same evaluation for all actions and store them in the action_values array, from which we derive the greedy policy. **Be careful when indexing into the value matrix**: `values[y, x]` stores the value of state (x, y)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oj97dbrea5cz"
      },
      "source": [
        "### Exercise 2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z7PNDZT-bsql"
      },
      "source": [
        "# Policy Iteration \n",
        "def policy_evaluation(grid, values, policy, gamma):\n",
        "    while True:\n",
        "        eps = 0\n",
        "        for y in range(grid.h):\n",
        "            for x in range(grid.w):\n",
        "                v = values[y, x]\n",
        "                ################################################################\n",
        "                # Fill in missing code below (...),\n",
        "                # then remove or comment the line below to test your function\n",
        "                raise NotImplementedError(\"Policy evaluation\")\n",
        "                ################################################################\n",
        "                new_v = ...\n",
        "                \n",
        "                values[y, x] = new_v\n",
        "                eps = max(eps, abs(new_v - v))\n",
        "        if eps < 0.0001: \n",
        "            break\n",
        "\n",
        "def policy_improvement(grid, values, policy, gamma):\n",
        "    converged = True\n",
        "    for y in range(grid.h):\n",
        "        for x in range(grid.w):\n",
        "            old_action = policy[y, x]\n",
        "            action_values = np.zeros(grid.n_actions, dtype=np.float)\n",
        "            ####################################################################\n",
        "            # Fill in missing code below (...),\n",
        "            # then remove or comment the line below to test your function\n",
        "            raise NotImplementedError(\"Policy improvement\")\n",
        "            ####################################################################\n",
        "            for action in range(...):\n",
        "                action_values[action] = ...\n",
        "            policy[y, x] = ...\n",
        "\n",
        "            if old_action != policy[y, x]:\n",
        "                converged = False\n",
        "    return converged\n",
        "\n",
        "def policy_iteration(grid, gamma=0.9):\n",
        "    policy = np.random.choice(grid.n_actions, (grid.h, grid.w)).astype(int)\n",
        "    values = np.zeros_like(grid.rew_grid)\n",
        "    converged = False\n",
        "    while not converged:\n",
        "        print(\"running policy evaluation\")\n",
        "        policy_evaluation(grid, values, policy, gamma)\n",
        "        print(\"running policy improvement\")\n",
        "        converged = policy_improvement(grid, values, policy, gamma)\n",
        "    return values, policy\n",
        "\n",
        "# # Uncomment to test\n",
        "# grid = get_book_grid()\n",
        "# values, policy = policy_iteration(grid)\n",
        "# grid.render_values(values)\n",
        "# grid.render_policy(policy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IsvHbqFJt4aA"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "# Policy Iteration \n",
        "def policy_evaluation(grid, values, policy, gamma):\n",
        "    while True:\n",
        "        eps = 0\n",
        "        for y in range(grid.h):\n",
        "            for x in range(grid.w):\n",
        "                v = values[y, x]\n",
        "                (new_x, new_y), reward = grid.get_transition((x, y), policy[y, x])\n",
        "                new_v = reward + gamma * values[new_y, new_x]\n",
        "                values[y, x] = new_v\n",
        "                eps = max(eps, abs(new_v - v))\n",
        "        if eps < 0.0001: \n",
        "            break\n",
        "\n",
        "def policy_improvement(grid, values, policy, gamma):\n",
        "    converged = True\n",
        "    for y in range(grid.h):\n",
        "        for x in range(grid.w):\n",
        "            old_action = policy[y, x]\n",
        "            action_values = np.zeros(grid.n_actions, dtype=np.float)\n",
        "            for action in range(grid.n_actions):\n",
        "                (new_x, new_y), reward = grid.get_transition((x, y), action)\n",
        "                action_values[action] = reward + gamma * values[new_y, new_x]\n",
        "            policy[y, x] = np.argmax(action_values)\n",
        "            if old_action != policy[y, x]:\n",
        "                converged = False\n",
        "    return converged\n",
        "\n",
        "def policy_iteration(grid, gamma=0.9):\n",
        "    policy = np.random.choice(grid.n_actions, (grid.h, grid.w)).astype(int)\n",
        "    values = np.zeros_like(grid.rew_grid)\n",
        "    converged = False\n",
        "    while not converged:\n",
        "        print(\"running policy evaluation\")\n",
        "        policy_evaluation(grid, values, policy, gamma)\n",
        "        print(\"running policy improvement\")\n",
        "        converged = policy_improvement(grid, values, policy, gamma)\n",
        "    return values, policy\n",
        "\n",
        "\n",
        "grid = get_book_grid()\n",
        "values, policy = policy_iteration(grid)\n",
        "grid.render_values(values)\n",
        "grid.render_policy(policy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v2trCXzobQfV"
      },
      "source": [
        "Construct the path from the policy visualization and see that following the policy from the initial state indeed leads to terminal state. Now change $\\gamma$ to 1.0 and rerun the code. Does policy iteration still converge? Why are we stuck on policy evaluation? (This is a brain-teaser, so don't spend too much time on it, and don't let the code run for too long.)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "eAiaKleibZcA"
      },
      "source": [
        "convergence = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8N0oqYEdrBiu"
      },
      "source": [
        "# to_remove solution\n",
        "# Policy iteration fails to converge because we're stuck at the first policy \n",
        "# evaluation step. The initial policy is all zeros, which means we take action \n",
        "# \"up\" at all states. Consider what happens at the top left corner. Going up at \n",
        "# this state will lead to a reward of -1 and transitioning back to itself. Since \n",
        "# gamma = 1, we have V(s) = V(s) - 1, which never converges. This can be avoided \n",
        "# by setting gamma < 1. The evaluation will halt when V(s) = gamma * V(s) - 1."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MwFYqD3kjE6h"
      },
      "source": [
        "## Section 4.2: Value Iteration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SDDeaHezfMve"
      },
      "source": [
        "Value iteration can be thought of as a simplification of policy iteration, where we effectively combine the two steps in policy iteration into one. We still iterate over all states, but in each iteration the value update becomes $$V(s) \\leftarrow \\max_a R(s, a) + \\gamma\\sum_{s'}P(s'|s, a)V(s') $$ So instead of computing the state value and then selecting the greedy action, we directly store the maximum state-action value. This obviates the need to maintain an explicit policy. After the value matrix has converged, we can back out the optimal policy by taking the argmax, same as what we did in policy improvement.\n",
        "\n",
        "Now it's your turn to implement the value iteration algorithm. You need to fill in the new update rule, and copy your code from policy improvment to reconstruct the optimal policy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epKCgDH5a89u"
      },
      "source": [
        "### Exercise 3"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ckLBXw2AtqfR"
      },
      "source": [
        "# Value Iteration\n",
        "\n",
        "def value_iteration(grid, gamma=0.9):\n",
        "    V = np.zeros_like(grid.rew_grid)\n",
        "    while True:\n",
        "        eps = 0\n",
        "        for y in range(grid.h):\n",
        "            for x in range(grid.w):\n",
        "                v = values[y, x]\n",
        "                action_values = np.zeros(grid.n_actions)\n",
        "                ################################################################\n",
        "                # Fill in missing code below (...),\n",
        "                # then remove or comment the line below to test your function\n",
        "                raise NotImplementedError(\"Value iteration\")\n",
        "                ################################################################\n",
        "                for action in range(...):\n",
        "                    action_values[action] = ... \n",
        "                new_v = ...\n",
        "                \n",
        "                values[y, x] = new_v\n",
        "                eps = max(eps, abs(new_v - v))\n",
        "        if eps < 0.0001:\n",
        "            break\n",
        "\n",
        "    # Create greedy policy from values\n",
        "    policy = np.zeros_like(grid.rew_grid).astype(int)\n",
        "    for y in range(grid.h):\n",
        "        for x in range(grid.w):\n",
        "            action_values = np.zeros(grid.n_actions)\n",
        "            ####################################################################\n",
        "            # Copy your solution for policy improvement here\n",
        "            raise NotImplementedError(\"Value iteration policy\")\n",
        "            #################################################################### \n",
        "            for action in range(...):\n",
        "                action_values[action] = ...\n",
        "            policy[y, x] = ...\n",
        "\n",
        "    return values, policy\n",
        "\n",
        "# # Uncomment to test\n",
        "# grid = get_book_grid()\n",
        "# values, policy = value_iteration(grid)\n",
        "# grid.render_values(values)\n",
        "# grid.render_policy(policy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LmXBn399t71Q"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "# Value Iteration\n",
        "def value_iteration(grid, gamma=0.9):\n",
        "    V = np.zeros_like(grid.rew_grid)\n",
        "    while True:\n",
        "        eps = 0\n",
        "        for y in range(grid.h):\n",
        "            for x in range(grid.w):\n",
        "                v = values[y, x]\n",
        "                action_values = np.zeros(grid.n_actions)\n",
        "                for action in range(grid.n_actions):\n",
        "                    (nx, ny), reward = grid.get_transition((x, y), action)\n",
        "                    action_values[action] = reward + gamma * values[ny, nx] \n",
        "                new_v = np.max(action_values)\n",
        "                values[y, x] = new_v\n",
        "                eps = max(eps, abs(new_v - v))\n",
        "        if eps < 0.0001:\n",
        "            break\n",
        "\n",
        "    # Create greedy policy from values\n",
        "    policy = np.zeros_like(grid.rew_grid).astype(int)\n",
        "    for y in range(grid.h):\n",
        "        for x in range(grid.w):\n",
        "            action_values = np.zeros(grid.n_actions)\n",
        "            for action in range(grid.n_actions):\n",
        "                (nx, ny), reward = grid.get_transition((x, y), action)\n",
        "                action_values[action] = reward + gamma * values[ny, nx] \n",
        "            policy[y, x] = np.argmax(action_values)\n",
        "\n",
        "    return values, policy\n",
        "\n",
        "grid = get_book_grid()\n",
        "values, policy = value_iteration(grid)\n",
        "grid.render_values(values)\n",
        "grid.render_policy(policy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NqEGcKG7MNBN"
      },
      "source": [
        "---\n",
        "# Section 5: Temporal Difference (TD) Learning\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WWz3Nvcs4pSa",
        "cellView": "form"
      },
      "source": [
        "#@title Video : TD and Q Learning\n",
        "try: t4;\n",
        "except NameError: t4=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"rCk_hvwZ6iA\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D0cMYLme3-3U"
      },
      "source": [
        "## Section 5.1 Q-learning\n",
        "\n",
        "Up until now we have assumed full access to the transitions of an environment. But in a typical reinforcement learning problem the dynamics is unknown. So how do we solve it? One way is to learn to approximate the dynamics using a function approximator (e.g. a neural net) and then apply dynamic programming or trajectory optimization. This is called model-based reinforcement learning, which we will cover next week. In this tutorial, we will study algorithms in the model-free regime. Specifically, we will investigate **Temporal Difference (TD) learning**.\n",
        "\n",
        "The idea behind TD learning is to use $V(s_{t+1})$ as an imperfect proxy for the true value (Monte Carlo bootstrapping), and obtain a generalized equation to calculate the TD error:\n",
        "$$\\delta_t = r_{t+1} + \\gamma V(s_{t+1}) - V(s_t)$$\n",
        "\n",
        "The expression $r_{t+1} + \\gamma V(s_{t+1})$ is also called the TD target. We can then update the value using a learning rate $\\alpha$.\n",
        "$$ V(s_t) \\leftarrow V(s_t) + \\alpha \\delta_t$$\n",
        "\n",
        "**Q-learning** is an instantiation of TD learning, where the TD error is $$\\delta_t = R(s_t, a_t) + \\gamma \\max_{a} Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t)$$ and the full update rule is $$Q(s_t,a_t) \\leftarrow Q(s_t, a_t) + \\alpha \\left(R(s_t, a_t) + \\gamma \\max_{a} Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t)\\right)$$\n",
        "\n",
        "Because of the max operator used to select the optimal Q-value in the TD target, Q-learning directly estimates the optimal action value, i.e. the cumulative future reward that would be obtained if the agent behaved optimally, regardless of the policy currently followed by the agent. For this reason, Q-learning is referred to as an **off-policy** method. A sketch of the Q-learning algorithm is as follows:\n",
        "\n",
        "```\n",
        "for n episodes:\n",
        "    for T steps:\n",
        "        Select an action a_t using some policy derived from the current Q-values\n",
        "        Execute a_t in the environment to get reward r and next state s_{t+1}\n",
        "        Update Q(s_t, a_t) using (s_t, a_t, r, s_{t+1})\n",
        "```\n",
        "\n",
        "A remaining question is, how do we select an action base on the current Q-values? If the approximated Q-values are very bad, then greedily following the argmax policy may cause the agent to get stuck in some bad states. Thus, we instead adopt an **epsilon-greedy policy**, where we choose the argmax action with probability $(1-\\epsilon)$ and take a random action otherwise. This relates to an important concept in reinforcement learning, namely exploration vs. exploitation. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EehisZexviwX"
      },
      "source": [
        "# Epsilon-greedy policy\n",
        "def epsilon_greedy(q_values, epsilon):\n",
        "    if np.random.random() > epsilon:\n",
        "        action = np.argmax(q_values)\n",
        "    else:\n",
        "        action = np.random.choice(len(q_values))\n",
        "    return action\n",
        "\n",
        "# General TD learning algorithm\n",
        "def learn_gridworld(env, backup_rule, params, max_steps, n_episodes):\n",
        "    values = np.zeros((env.h, env.w, env.n_actions))\n",
        "    episode_actions = []\n",
        "    episode_rewards = np.zeros(n_episodes)\n",
        "\n",
        "    for episode in tqdm(range(n_episodes)):\n",
        "        env.reset()\n",
        "        total_reward = 0\n",
        "        action_list = []\n",
        "\n",
        "        for t in range(max_steps):\n",
        "            state = env.state\n",
        "            # Select action from epsilon-greedy policy\n",
        "            action = epsilon_greedy(values[state[1], state[0]], params['epsilon'])\n",
        "            action_list.append(action)\n",
        "            # Execute action\n",
        "            next_state, reward, done = env.step(action)\n",
        "            # Update values\n",
        "            values = backup_rule(state, action, reward, next_state, values, params)\n",
        "            total_reward += reward\n",
        "            if done:\n",
        "                break\n",
        "        \n",
        "        episode_actions.append(action_list)\n",
        "        episode_rewards[episode] = total_reward\n",
        "\n",
        "    return values, episode_rewards"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpnKs95GbFdC"
      },
      "source": [
        "### Exercise 4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A6QbhpiSiLub"
      },
      "source": [
        "In this exercise, you will implement the update rule for Q-learning and test it on the Cliff World environment, where the agent needs to navigate to the other side of the cliff without falling off. You need to fill in the code for computing the TD error and updating the values matrix."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lTZGFohOOvYD"
      },
      "source": [
        "# Q-Learning\n",
        "def q_learning_backup(state, action, reward, next_state, values, params):\n",
        "    '''\n",
        "    Compute a new set of q-values using the q-learning update rule.\n",
        "    Args:\n",
        "        state (tuple): s_t, a tuple of xy coordinates.\n",
        "        action (int): a_t, an integer from {0, 1, 2, 3}.\n",
        "        reward (float): the reward of executing a_t at s_t.\n",
        "        next_state (tuple): s_{t+1}, a tuple of xy coordinates.\n",
        "        values (ndarray): an (h, w, 4) numpy array of q-values. values[y, x, a] \n",
        "                          stores the value of executing action a at state (x, y).\n",
        "        params (dict): a dictionary of parameters.\n",
        "\n",
        "    Returns:\n",
        "        ndarray: the updated q-values.\n",
        "    '''\n",
        "    x, y = state\n",
        "    nx, ny = next_state\n",
        "    gamma = params['gamma']\n",
        "    alpha = params['alpha']\n",
        "\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Q-learning\")\n",
        "    ####################################################################\n",
        "    q = ...\n",
        "    max_next_q = ...\n",
        "\n",
        "    # Compute TD error using q and max_next_q\n",
        "    td_error = ...\n",
        "    values[y, x, action] = ...\n",
        "\n",
        "    return values\n",
        "\n",
        "# # Uncomment to test\n",
        "# env = get_cliff_walk()\n",
        "# params = {'gamma': 1.0, 'alpha': 0.1 , 'epsilon': 0.1}\n",
        "# max_steps = 1000\n",
        "# n_episodes = 500\n",
        "# q_values, episode_rewards = learn_gridworld(env, q_learning_backup, params, max_steps, n_episodes)\n",
        "\n",
        "# plot_episode_rewards(episode_rewards)\n",
        "# env.render_policy(np.argmax(q_values, axis=2))\n",
        "# env.render_q_values(q_values)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39eQr6zMuB4e"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "# Q-learning\n",
        "def q_learning_backup(state, action, reward, next_state, values, params):\n",
        "    x, y = state\n",
        "    nx, ny = next_state\n",
        "    gamma = params['gamma']\n",
        "    alpha = params['alpha']\n",
        "    \n",
        "    q = values[y, x, action]\n",
        "    max_next_q = np.max(values[ny, nx])\n",
        "\n",
        "    td_error = reward + (gamma * max_next_q - q)\n",
        "    values[y, x, action] = q + alpha * td_error\n",
        "\n",
        "    return values\n",
        "\n",
        "env = get_cliff_walk()\n",
        "params = {'gamma': 1.0, 'alpha': 0.1 , 'epsilon': 0.1}\n",
        "max_steps = 1000\n",
        "n_episodes = 500\n",
        "q_values, episode_rewards = learn_gridworld(env, q_learning_backup, params, max_steps, n_episodes)\n",
        "\n",
        "plot_episode_rewards(episode_rewards)\n",
        "env.render_policy(np.argmax(q_values, axis=2))\n",
        "env.render_q_values(q_values)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7_3P3ZRB3896"
      },
      "source": [
        "## Section 5.2: SARSA"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HM7fOpXcUbIa"
      },
      "source": [
        "An alternative to Q-learning, the SARSA algorithm also estimates action values. However, rather than estimating the optimal (off-policy) values, SARSA estimates the **on-policy** action values, i.e. the cumulative future reward that would be obtained if the agent behaved according to its current beliefs.\n",
        "\n",
        "\\begin{align}\n",
        "Q(s_t,a_t) \\leftarrow Q(s_t,a_t) + \\alpha \\big(R(s_t, a_t) + \\gamma Q(s_{t+1}, \\pi(s_{t+1})) - Q(s_t,a_t)\\big)\n",
        "\\end{align}\n",
        "\n",
        "In fact, you will notices that the *only* difference between Q-learning and SARSA is the TD target calculation uses the policy to select the next action (in our case epsilon-greedy) rather than using the action that maximizes the Q-value. You do not need to implement the SARSA algorithm. Run the following code cell and compare with Q-learning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lP3OrNyRt-K3"
      },
      "source": [
        "# SARSA\n",
        "def sarsa_backup(state, action, reward, next_state, values, params):\n",
        "    '''\n",
        "    Compute a new set of q-values using the SARSA update rule.\n",
        "    Args:\n",
        "        state (tuple): s_t, a tuple of xy coordinates.\n",
        "        action (int): a_t, an integer from {0, 1, 2, 3}.\n",
        "        reward (float): the reward of executing a_t at s_t.\n",
        "        next_state (tuple): s_{t+1}, a tuple of xy coordinates.\n",
        "        values (ndarray): an (h, w, 4) numpy array of q-values. values[y, x, a] \n",
        "                          stores the value of executing action a at state (x, y).\n",
        "        params (dict): a dictionary of parameters.\n",
        "\n",
        "    Returns:\n",
        "        ndarray: the updated q-values.\n",
        "    '''\n",
        "    x, y = state\n",
        "    nx, ny = next_state\n",
        "    gamma = params['gamma']\n",
        "    alpha = params['alpha']\n",
        "\n",
        "    q = values[y, x, action]\n",
        "    # Obtain on-policy action \n",
        "    policy_action = epsilon_greedy(values[ny, nx], params['epsilon'])\n",
        "    next_q = values[ny, nx, policy_action]\n",
        "\n",
        "    # Compute TD error using q and max_next_q\n",
        "    td_error = reward + (gamma * next_q - q)\n",
        "    values[y, x, action] = q + alpha * td_error\n",
        "\n",
        "    return values\n",
        "\n",
        "env = get_cliff_walk()\n",
        "params = {'gamma': 1.0, 'alpha': 0.1 , 'epsilon': 0.1}\n",
        "max_steps = 1000\n",
        "n_episodes = 500\n",
        "q_values, episode_rewards = learn_gridworld(env, sarsa_backup, params, max_steps, n_episodes)\n",
        "\n",
        "plot_episode_rewards(episode_rewards)\n",
        "env.render_policy(np.argmax(q_values, axis=2))\n",
        "env.render_q_values(q_values)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wzov7ZL54vYc"
      },
      "source": [
        "Compare the reward plots and policies of Q-learning and SARSA. Do they take the same path to reach the terminal state? Why does one look more conservative than the other?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "dZYizmriwniE"
      },
      "source": [
        "q_vs_sarsa = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3wBq9z5cXfBD"
      },
      "source": [
        "# to_remove solution\n",
        "# SARSA appears to be more conservative than Q-learning. The path taken by SARSA\n",
        "# is one step away from the cliff, whereas that taken by Q-learning is right on\n",
        "# the edge. We also see this behavior from the higher variance in the reward plot \n",
        "# of Q-learning. The reason is that SARSA constructs its TD target by taking one \n",
        "# more on-policy action. Consider updating the value of action \"right\" at an edge \n",
        "# cell. The TD target for Q-learning involves the value of the next greedy action, \n",
        "# which will converge to \"right\" after some iterations. The TD target for SARSA \n",
        "# involves the value of the next epsilon-greedy action, which always has a chance \n",
        "# of being \"down\" aka \"falling off the cliff.\" Thus, SARSA assign a lower values \n",
        "# to action \"right\" for edge cells."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4S4tPIYg0TeB"
      },
      "source": [
        "## Section 5.3 (Optional): Try your own grid\n",
        "\n",
        "If time allows, feel free to try Q-learning or SARSA on one of the other pre-defined grids or a Gridworld of your own creation. Discuss your findings with your pod."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HpkSH_BK49-b"
      },
      "source": [
        "---\n",
        "# Wrap-up and foreshadowing"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4_hSrDhG5IIv",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Wrap-up\n",
        "try: t5;\n",
        "except NameError: t5=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"oJo0jb_h2sM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SkggStBmyOnQ",
        "cellView": "form"
      },
      "source": [
        "import time\n",
        "import numpy as np\n",
        "import urllib.parse\n",
        "from IPython.display import IFrame\n",
        "\n",
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "\n",
        "\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefill_fields = {}\n",
        "  for key in fields:\n",
        "      new_key = 'prefill_' + key\n",
        "      prefill_fields[new_key] = fields[key]\n",
        "  prefills = urllib.parse.urlencode(prefill_fields)\n",
        "  src = src + prefills\n",
        "  return src\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "try: t4;\n",
        "except NameError: t4 = time.time()\n",
        "try: t5;\n",
        "except NameError: t5 = time.time()\n",
        "try: t6;\n",
        "except NameError: t6 = time.time()\n",
        "\n",
        "#autofill fields if they are not present\n",
        "#a missing pennkey and pod will result in an Airtable warning\n",
        "#which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: learning_from_previous_week;\n",
        "except NameError: learning_from_previous_week = \"\"\n",
        "try: MDP_example;\n",
        "except NameError: MDP_example = \"\"\n",
        "try: life_discount;\n",
        "except NameError: life_discount = \"\"\n",
        "try: convergence;\n",
        "except NameError: convergence = \"\"\n",
        "try: q_vs_sarsa;\n",
        "except NameError: q_vs_sarsa = \"\"\n",
        "\n",
        "\n",
        "times = np.array([t1,t2,t3,t4,t5,t6])-t0\n",
        "\n",
        "fields = {\"pennkey\": my_pennkey,\n",
        "          \"pod\": my_pod,\n",
        "          \"learning_from_previous_week\": learning_from_previous_week,\n",
        "          \"MDP_example\": MDP_example,\n",
        "          \"life_discount\": life_discount,\n",
        "          \"convergence\": convergence,\n",
        "          \"q_vs_sarsa\": q_vs_sarsa,\n",
        "          \"cumulative_times\": times}\n",
        "\n",
        "src = \"https://airtable.com/embed/shrS0Ltpj30NO4Fr8?\"\n",
        "\n",
        "\n",
        "# now instead of the original source url, we do: src = prefill_form(src, fields)\n",
        "display(IFrame(src = prefill_form(src, fields), width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hPcTr2XWyehq"
      },
      "source": [
        "## Feedback\n",
        "How could this session have been better? How happy are you in your group? How do you feel right now?\n",
        "\n",
        "Feel free to use the embeded form below or use this link:\n",
        "<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://airtable.com/shrNSJ5ECXhNhsYss\">https://airtable.com/shrNSJ5ECXhNhsYss</a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "popba9fsyhYa"
      },
      "source": [
        "display(IFrame(src=\"https://airtable.com/embed/shrNSJ5ECXhNhsYss?backgroundColor=red\", width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}